from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain.chat_models import init_chat_model
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.prebuilt import ToolNode
from langchain_ollama import ChatOllama
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessage,
    HumanMessage,
)
import asyncio

async def main():
    # Initialize the model
    model = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False, base_url="http://127.0.0.1:11434/")

    # Set up MCP client
    client = MultiServerMCPClient(
        {
            "models": {
                # make sure you start your weather server on port 8000
                "url": "https://u-wanghh2000-model-tools-c1.space.opencsg.com/sse",
                # "transport": "streamable_http",
                "transport": "sse",
            }
        }
    )
    tools = await client.get_tools()

    # Bind tools to model
    model_with_tools = model.bind_tools(tools)

    # Create ToolNode
    tool_node = ToolNode(tools)

    def should_continue(state: MessagesState):
        messages = state["messages"]
        last_message = messages[-1]
        if last_message.tool_calls:
            return "tools"
        return END

    # Define call_model function
    async def call_model(state: MessagesState):
        messages = state["messages"]
        response = await model_with_tools.ainvoke(messages)
        return {"messages": [response]}

    # Build the graph
    builder = StateGraph(MessagesState)
    builder.add_node("call_model", call_model)
    builder.add_node("tools", tool_node)

    builder.add_edge(START, "call_model")
    builder.add_conditional_edges(
        "call_model",
        should_continue,
    )
    builder.add_edge("tools", "call_model")

    # Compile the graph
    graph = builder.compile()

    # Test the graph
    # math_response = await graph.ainvoke(
    #     {"messages": [{"role": "user", "content": "what's (3 + 5) x 12?"}]}
    # )
    # weather_response = await graph.ainvoke(
    #     {"messages": [{"role": "user", "content": "what is the weather in nyc?"}]}
    # )
    models_response = await graph.ainvoke(
        {"messages": [HumanMessage(content="what are the most popular models?")]},
        # {"messages": [{"role": "user", "content": "what are the most popular models?"}]}
    )
    print(models_response["messages"][-1])


asyncio.run(main())
